/* Copyright (c) 2009, 2010, 2011 Damian Kmiecik <d0zoenator@gmail.com>
   All rights reserved.

   Redistribution and use in source and binary forms, with or without
   modification, are permitted provided that the following conditions are met:

   * Redistributions of source code must retain the above copyright
     notice, this list of conditions and the following disclaimer.
   * Redistributions in binary form must reproduce the above copyright
     notice, this list of conditions and the following disclaimer in
     the documentation and/or other materials provided with the
     distribution.
   * Neither the name of the copyright holders nor the names of
     contributors may be used to endorse or promote products derived
     from this software without specific prior written permission.

  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
  AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
  IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
  ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
  LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
  CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
  SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
  INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
  CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
  ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
  POSSIBILITY OF SUCH DAMAGE. */

#include <stdlib.h>
#include <avr/interrupt.h>
#include <avr/io.h>
#include <util/delay.h>
#include <avr/pgmspace.h>
#include <stdio.h>
#include "main.h"
#include "checksum.h"
#include "timer.h"
#include "stdout.h"
#include "enc28j60.h"
#include "ethernet.h"

#ifndef TCP_DEBUG
	#define TCP_DEBUG(...)
#endif

#ifndef TCP_APPS_RUNNER
	#error You must define callback function TCP_APPS_RUNNER! Reference to the manual first.
#endif
#ifndef UDP_APPS_RUNNER
	#error You must define callback function UDP_APPS_RUNNER! Reference to the manual first.
#endif

struct TCP_PORT_ITEM TCP_PORT_TABLE[MAX_APP_ENTRY+1] = {
		{0},
		{0},
		{0}
};
struct UDP_PORT_ITEM UDP_PORT_TABLE[MAX_APP_ENTRY] = {
		{0},
		{0},
		{0}
};
struct TCP_TABLE tcpTable[MAX_TCP_ENTRY+1];
struct ETH_UDP_HEADER *udp = (struct ETH_UDP_HEADER*)&ethBuffer[UDP_OFFSET];
struct ETH_TCP_HEADER *tcp = (struct ETH_TCP_HEADER*)&ethBuffer[TCP_OFFSET];
struct TCP_TABLE *tcp_conn;
struct TCP_STATUS tcp_status;

#ifdef DEBUG_TCP
void TCPprintTable(void)
{
	printf_P(PSTR("_ TCP TABLE _\n"));
	for (uint8_t i = 0; i < MAX_TCP_ENTRY; i++)
	{
		//if (tcpTable[i].ip != 0)
			printf_P(
				PSTR("%u.%u.%u.%u dest %u src %u for %us\n"),
				((uint8_t*)&(tcpTable[i].ip))[0],
				((uint8_t*)&(tcpTable[i].ip))[1],
				((uint8_t*)&(tcpTable[i].ip))[2],
				((uint8_t*)&(tcpTable[i].ip))[3],
				tcpTable[i].dest_port,
				tcpTable[i].src_port,
				tcpTable[i].time
			);
	}
	printf_P(PSTR("_ TCP TABLE END _\n"));
}
#endif

uint8_t TCPportSearch(const uint16_t destPort)
{
	for (uint8_t portIndex = 0; portIndex < MAX_APP_ENTRY; portIndex++)
		if (TCP_PORT_TABLE[portIndex].port == destPort)
			return portIndex;
	return MAX_APP_ENTRY;
}

void TCPtimerDo(void)
{
	for (uint8_t tcpIndex = 0; tcpIndex < MAX_TCP_ENTRY; tcpIndex++)
	{
		if (tcpTable[tcpIndex].ip != 0)
		{
			tcp_conn = &tcpTable[tcpIndex];
			if (tcp_conn->time == 0)
			{
				if (tcp_conn->send)
				{
					if (tcp_conn->err--)
					{
						TCP_DEBUG("TCP %hu: RETRANSMISSION\n", tcpIndex);
						tcp_status.status = TCP_REXMIT;
					} else {
						TCP_DEBUG("TCP %hu: RETRANSMISSION TIMEDOUT\n", tcpIndex);
						tcp_status.status = TCP_TIMEDOUT;
					}
				} else {
					TCP_DEBUG("TCP %hu: TIMEDOUT\n", tcpIndex);
					tcp_status.status = TCP_TIMEDOUT;
				}
				TCP_APPS_RUNNER();
			} else {
				tcp_conn->time--;
			}
			if (!(tcp_status.status & TCP_TIMEDOUT))
			{
				tcp_status.status = TCP_POOL;
				TCP_APPS_RUNNER();
			}
		}
	}
}

void TCPindexDel(void)
{
	tcp_conn->ip = 0;
	tcp_conn->ack_counter = 0;
	tcp_conn->dest_port = 0;
	tcp_conn->first_ack = 0;
	tcp_conn->flags = 0;
	tcp_conn->ip = 0;
	tcp_conn->seq_counter = 0;
	tcp_conn->src_port = 0;
	tcp_conn->send = 0;
	tcp_conn->time = 0;
}

void TCPappChangePort(const port_t port_old, const port_t port_new)
{
	uint8_t i;
	for (i = 0; i < MAX_APP_ENTRY; i++)
	{
		if (TCP_PORT_TABLE[i].port == port_old)
		{
			TCP_DEBUG("Note: TCP change port %u to %u\n",
				htons(port_old),
				htons(port_new)
			);
			TCP_PORT_TABLE[i].port = port_new;
			return;
		}
	}
}

void TCPentryAdd(void)
{
	uint8_t index;
	uint32_t result32;
	// Update
	for (index = 0; index < MAX_TCP_ENTRY; index++)
	{
		if ((tcpTable[index].ip == ip->sourceIp) &&
				(tcpTable[index].src_port == tcp->sourcePort))
		{
			tcp_conn = &tcpTable[index];
			tcp_conn->ack_counter = tcp->ackNum;
			tcp_conn->seq_counter = tcp->seqNum;
			tcp_conn->flags = tcp->flags;
			if (tcp_conn->time != TCP_TIME_OFF)
				tcp_conn->time = TCP_MAX_TIME;
			result32 = htons(ip->totalLen) -
					ETH_IP_HEADER_LEN -
					((tcp->offset & 0xF0) >> 2);
			result32 = result32 + htons32(tcp->seqNum);
			tcp_conn->seq_counter = htons32(result32);
			return;
		}
	}
	// Add new
	for (index = 0; index < MAX_TCP_ENTRY; index++)
	{
		if (tcpTable[index].ip == 0)
		{
			tcp_conn = &tcpTable[index];
			tcp_conn->ip = ip->sourceIp;
			tcp_conn->src_port = tcp->sourcePort;
			tcp_conn->dest_port = tcp->destPort;
			tcp_conn->ack_counter = tcp->ackNum;
			tcp_conn->seq_counter = tcp->seqNum;
			tcp_conn->flags = tcp->flags;
			tcp_conn->time = TCP_MAX_TIME;
			tcp_conn->send = 0;
			return;
		}
	}
	// No free space
	TCP_DEBUG("TCP Stack full\n", NULL);
	tcp_conn = NULL;
}

void TCPentrySearch(const uint32_t srcIp, const uint16_t srcPort)
{
	// Search
	for (uint8_t index = 0; index < MAX_TCP_ENTRY; index++)
	{
		if ((tcpTable[index].ip == srcIp) && (tcpTable[index].src_port
				== srcPort))
		{
			tcp_conn = &tcpTable[index];
			return;
		}
	}
	tcp_conn = NULL;
}

void TCPmakeNewPacket(uint16_t len)
{
	uint16_t result16;
	uint16_t bufferLen;
	uint32_t result32;
	tcp->sourcePort = tcp_conn->dest_port;
	tcp->destPort = tcp_conn->src_port;
	tcp->urgentPtr = 0;
	tcp->window = HTONS(MAX_TCP_WINDOW);
	tcp->offset = 0x50;
	tcp->flags = tcp_conn->flags;
	result32 = htons32(tcp_conn->seq_counter);
	if (tcp_conn->flags & TCP_SYN_FLAG) {
		tcp_conn->ack_counter = 0xe6acb148;
		result32++;
		ethBuffer[TCP_DATA] = 2;
		ethBuffer[TCP_DATA + 1] = 4;
		ethBuffer[TCP_DATA + 2] = (MAX_TCP_WINDOW >> 8) & 0xff;
		ethBuffer[TCP_DATA + 3] = MAX_TCP_WINDOW & 0xff;
		len = 0x04;
		tcp->offset = 0x60;
	}
	tcp->ackNum = htons32(result32);
	tcp->seqNum = tcp_conn->ack_counter;
	bufferLen = ETH_IP_HEADER_LEN + ETH_TCP_HEADER_LEN + len;
	ip->totalLen = htons(bufferLen);
	bufferLen += ETH_ETHERNET_HEADER_LEN;
	ip->protocol = IP_PROTOCOL_TCP;
	IPmakeHeader(tcp_conn->ip);
	tcp->checksum = 0;
	result16 = htons(ip->totalLen) + 8;
	result16 = result16 - ((ip->version & 0x0F) << 2);
	result32 = result16 - 2;
	result16 = checksum((&ip->version + 12), result16, result32);
	tcp->checksum = htons(result16);
	// Send packet
	enc28j60SendPacket(bufferLen, ethBuffer);
}

void TCPpacketProcess(void)
{
	uint32_t result32;
	uint8_t portIndex;
	// Search port in table
	portIndex = TCPportSearch(tcp->destPort);
	if (portIndex == MAX_APP_ENTRY)
	{
		TCP_DEBUG("TCP: port %u is closed\n", htons(tcp->destPort));
		return;
	}
	// Check if SYN+ACK requested and received
	if ((tcp->flags & TCP_SYN_FLAG) && (tcp->flags & TCP_ACK_FLAG))
	{
		TCPentryAdd();
		if (tcp_conn == NULL)
		{
			return;
		}
		tcp_conn->time = TCP_MAX_TIME;
		tcp_conn->seq_counter = htons32(
				htons32(tcp_conn->seq_counter) + 1
		);
		tcp_conn->flags = TCP_ACK_FLAG;
		TCPmakeNewPacket(0);
		tcp_conn->first_ack = 1;
		tcp_status.status = TCP_CONNECTED;
		TCP_APPS_RUNNER();
		return;
	}
	// Check if SYN requested and received
	if (tcp->flags == TCP_SYN_FLAG)
	{
		TCPentryAdd();
		if (tcp_conn == NULL)
		{
			return;
		}
		TCP_DEBUG("TCP: Connection estabilished %01hu.%01hu.%01hu.%01hu:%u\n",
			((uint8_t*) &(tcp_conn->ip))[0],
			((uint8_t*) &(tcp_conn->ip))[1],
			((uint8_t*) &(tcp_conn->ip))[2],
			((uint8_t*) &(tcp_conn->ip))[3],
			tcp_conn->src_port
		);
		tcp_conn->flags = TCP_SYN_FLAG | TCP_ACK_FLAG;
		TCPmakeNewPacket(0);
		return;
	}
	// Search
	TCPentrySearch(ip->sourceIp, tcp->sourcePort);
	// TCP not found
	if (tcp_conn == NULL)
	{
		DEBUG("TCP: Entry not found!\n", NULL);
		if ((tcp->flags & TCP_FIN_FLAG) || (tcp->flags & TCP_RST_FLAG))
		{
			// XXX: TEMPORARY TCP ENTRY!!!
			tcp_conn = &tcpTable[MAX_TCP_ENTRY];
			TCPentryAdd();
			result32 = htons32(tcp_conn->seq_counter) + 1;
			tcp_conn->seq_counter = htons32(result32);
			if (tcp_conn->flags & TCP_FIN_FLAG)
			{
				tcp_conn->flags = TCP_ACK_FLAG;
				TCPmakeNewPacket(0);
			}
			TCPindexDel();
		}
		return;
	}
	// Refresh TCP table
	TCPentryAdd();
	// Check if connection closeing
	if ((tcp->flags & TCP_FIN_FLAG) || (tcp->flags & TCP_RST_FLAG))
	{
		result32 = htons32(tcp_conn->seq_counter) + 1;
		tcp_conn->seq_counter = htons32(result32);
		if (tcp_conn->flags & TCP_FIN_FLAG)
		{
			tcp_conn->flags = TCP_FIN_FLAG|TCP_ACK_FLAG;
			TCPmakeNewPacket(0);
			tcp_status.status = TCP_CLOSED;
		} else {
			tcp_status.status = TCP_ABORTED;
		}
		// End the application
		TCP_APPS_RUNNER();
		TCPindexDel();
		return;
	}
	// Check if PSH+ACK flag
	if ((tcp_conn->flags & TCP_PSH_FLAG) &&
			(tcp_conn->flags & TCP_ACK_FLAG))
	{
		tcp_conn->flags =  TCP_ACK_FLAG;
		TCPmakeNewPacket(0);
		tcp_status.status = TCP_NEWDATA;
		TCP_APPS_RUNNER();
		return;
	}
	// Check if FIRST-ACK received
	if((tcp_conn->flags & TCP_ACK_FLAG) &&
			(tcp_conn->first_ack == 0))
	{
		tcp_conn->first_ack = 1;
		tcp_status.status = TCP_CONNECTED;
		TCP_APPS_RUNNER();
		return;
	}
	// Check if ACK flag; data received
	if((tcp_conn->flags & TCP_ACK_FLAG) &&
			(tcp_conn->first_ack == 1))
	{
		if (tcp_conn->send)
		{
			tcp_conn->send = 0;
			tcp_status.status = TCP_ACKED;
			TCP_APPS_RUNNER();
			return;
		}
		tcp_conn->flags =  TCP_ACK_FLAG;
		TCPmakeNewPacket(0);
		tcp_status.status = TCP_NEWDATA;
		TCP_APPS_RUNNER();
		return;
	}
}

void TCPconnect(const ip_t destIp, const port_t destPort)
{
	// Serach free port
	static uint16_t tcpPort = 0xFFFF;
	if (tcpPort < 4000)
		tcpPort = 0xFFFF;
	do {
		tcpPort--;
	} while (TCPportSearch(htons(tcpPort)) != MAX_APP_ENTRY);
	//
	uint8_t tcpIndex;
	for (tcpIndex = 0; tcpIndex < MAX_TCP_ENTRY; tcpIndex++)
	{
		if (tcpTable[tcpIndex].ip == 0)
		{
			tcp_conn = &tcpTable[tcpIndex];
			tcp_conn->ip = destIp;
			tcp_conn->src_port = destPort;
			tcp_conn->dest_port = htons(tcpPort);
			tcp_conn->ack_counter = 1234;
			tcp_conn->seq_counter = 2345;
			tcp_conn->time = TCP_MAX_TIME;
			break;
		}
	}
	if (tcpIndex == MAX_TCP_ENTRY)
	{
		tcp_conn = NULL;
		TCP_DEBUG("ERROR: Server busy\n", NULL);
		return;
	}
	TCP_DEBUG("TCP: Outcomming connection on port %u\n", tcpPort);
	tcp_conn->flags = TCP_SYN_FLAG;
	TCPmakeNewPacket(0);
}

void TCPportClose(void)
{
	TCP_DEBUG("TCP: Port close\n", NULL);
	tcp_conn->flags = TCP_ACK_FLAG | TCP_FIN_FLAG;
	TCPmakeNewPacket(0);
}

void TCPsend(const uint16_t size)
{
	tcp_conn->flags = TCP_PSH_FLAG | TCP_ACK_FLAG;
	tcp_conn->send = 1;
	tcp_conn->time = TCP_REXMIT_TIME;
	tcp_conn->err = TCP_MAX_ERROR_CNT;
	TCPmakeNewPacket(size);
}

void TCPlisten(const port_t port)
{
	uint8_t i;
	for (i = 0; i < MAX_APP_ENTRY; i++)
	{
		if (TCP_PORT_TABLE[i].port == 0)
		{
			TCP_DEBUG("TCP: Listening on port %u\n", htons(port));
			TCP_PORT_TABLE[i].port = port;
			return;
		}
	}
}

void UDPlisten(const port_t port)
{
	for (uint8_t i = 0; i < MAX_APP_ENTRY; i++)
	{
		if (UDP_PORT_TABLE[i].port == 0)
		{
			TCP_DEBUG("Note: UDP open port %u\n", htons(port));
			UDP_PORT_TABLE[i].port = port;
			return;
		}
	}
}

void UDPmakeNewPacket(const uint16_t len, const port_t srcPort,
		const port_t dstPort, const ip_t ipAddr)
{
	uint16_t bufferLen;
	udp->checksum = 0;
	udp->length = htons(len + ETH_UDP_HEADER_LEN);
	udp->sourcePort = srcPort;
	udp->destPort = dstPort;
	bufferLen = ETH_IP_HEADER_LEN + ETH_UDP_HEADER_LEN + len;
	ip->totalLen = htons(bufferLen);
	bufferLen += ETH_ETHERNET_HEADER_LEN;
	ip->protocol = IP_PROTOCOL_UDP;
	IPmakeHeader(ipAddr);
	// Send packet
	enc28j60SendPacket(bufferLen, ethBuffer);
}

void UDPpacketProcess(void)
{
	// Search port in table
	for (uint8_t i = 0; i < MAX_APP_ENTRY; i++)
	{
		if (UDP_PORT_TABLE[i].port == udp->destPort)
		{
			UDP_APPS_RUNNER();
			return;
		}
	}
	TCP_DEBUG("UDP blocked port %u\n", htons(udp->destPort));
}
